# -*- coding: utf-8 -*-  
'''
自定义回调

Created on 2021年8月23日
@author: luoyi
'''
import pickle
import tensorflow as tf

import utils.conf as conf


#    n轮batch，和每轮epoch后自动保存LDA相关文件
class AutoSaveGSDMMCallable(tf.keras.callbacks.Callback):
    '''n轮batch，和m轮epoch后自动保存GSDMM相关文件'''
    def __init__(self,
                 n_batch=10,
                 m_epoch=1,
                 wn_topic_word_path=conf.GSDMM.get_wn_topic_word_path(),
                 prob_topic_word_path=conf.GSDMM.get_prob_topic_word_path()):
        self._n_batch = n_batch
        self._m_epoch = m_epoch
        
        self._wn_topic_word_path = wn_topic_word_path
        self._prob_topic_word_path = prob_topic_word_path
        
        #    先建目录
        conf.mkfiledir_ifnot_exises(self._wn_topic_word_path)
        conf.mkfiledir_ifnot_exises(self._prob_topic_word_path)
        pass

    #    判断n轮batch后
    def on_batch_end(self, batch, logs=None):
        if (self._n_batch > 0 and batch % self._n_batch == 0):
            self.save_lda()
        pass
    #    判断m轮epoch后保存lda
    def on_epoch_end(self, epoch, logs=None):
        if (self._m_epoch > 0 and epoch % self._m_epoch == 0):
            self.save_lda()
        pass
    
    #    保存lda相关参数
    def save_lda(self):
        #    取主题-词 相关的 计数 与 概率
        wn_topic_word = self.model._prob_layer.wn_topic_word()
        topic_word_prob = self.model._prob_layer.topic_word_prob()
        
        #    保存进文件
        with open(file=self._wn_topic_word_path, mode='wb') as fw: pickle.dump(wn_topic_word.numpy(), fw)
        with open(file=self._prob_topic_word_path, mode='wb') as fw: pickle.dump(topic_word_prob.numpy(), fw)
        pass
    
    pass
